#! /usr/bin/env python
# coding=utf-8


# 通过读取结果获得identity的cutoff 的档次
# 然后丛对应的mRNA中 提取序列 并进行翻译

import os
import shutil
import numpy as np
from Bio import SeqIO

from Bio import Seq


import re
import sys
import argparse



dis = '''
对制定物种的tblastn 结果 提取mRNA和pep的原始序列
tblastn_fetch_gene01_fetch_gene.py -p raw -m db -t tblastn_res -i target_lista
'''


parser = argparse.ArgumentParser(description=dis)





parser.add_argument('-p',
                help='必须给定，结果前缀')

parser.add_argument('-m',
                help='必须给定，存有mRNA fasta的路径 名称需要是 XXXXX.fa XXXXX为物种缩写')

parser.add_argument('-t',
                help='必须给定，tblastn 结果路径')

parser.add_argument('-i',
                help='必须给定，目标物种列表文件')




# 获取参数
args = parser.parse_args()



if not args.p or not args.m or not args.t or not args.i:
    parser.print_help()
    sys.exit()





target_file = args.i

tblastn_res = args.t

mRNA_file = args.m
outfile_prefix = args.p




outfile_mRNA = open(outfile_prefix+'_mRNA.fa','w')
outfile_pep = open(outfile_prefix+'_pep.fa','w')

def find_orfs_with_trans(seq, trans_table, min_protein_length):
    
    answer = []
    seq_len = len(seq)

    trans = seq.translate(trans_table)
    trans_len = len(trans)
    aa_start = 0
    aa_end = 0
    while aa_start < trans_len:
        aa_end = trans.find("*", aa_start)
        if aa_end == -1:
            aa_end = trans_len
        if aa_end - aa_start >= min_protein_length:
         
            start = aa_start * 3
            end = min(seq_len, aa_end * 3 + 3)

            answer.append((start, end, trans[aa_start:aa_end]))
        aa_start = aa_end + 1
    answer = sorted(answer,key=lambda x:-int(x[1]-x[0]))
    return answer





target_lista= []
with open(target_file) as fila:
    for i in fila:
        k = i.strip()
        if k!='':
            target_lista.append(k)

 
value_lista = {}
tblastn_item = {}
max_size = 0
min_size = 99999
for i in os.listdir(tblastn_res):
    if os.path.exists(tblastn_res+'/'+i):
        with open(tblastn_res+'/'+i) as fila:
            for k in fila:

                k = k.strip().split('\t')
                if len(k)>5:
                    size = int(k[7])-int(k[6])
                    if size > max_size:
                        max_size = size
                    if size < min_size:
                        min_size = size
                    value_lista[i] = float(k[2])/100*size
                    tblastn_item[i] = k



# cutoff 设计为 取一半的档次


value_lista_tmp = np.array(list(value_lista.values()))/max_size
cut_off = np.median(value_lista_tmp)*100//10*10/100 * max_size
print(value_lista_tmp)

# print(value_lista)


print('cutoff is %s' %cut_off)
if cut_off<40: # 相似度 小于40的直接排除
    print('cutoff too low, use 50')
    cut_off = 50
    #sys.exit()


spe_lista = [i for i in value_lista if value_lista[i] >= cut_off]


# print(spe_lista)




print('select spe is %s' % ','.join(spe_lista))




# 获取对应的pep 和mRNA 最终以 fasta 输出到制定文件

#ttt=  1

for i in spe_lista:

    scaf = i[1]
    if not os.path.exists(mRNA_file+'/'+i+'.fa'):
        print('错误 mRNA文件 %s.fa 不存在于 %s'%(i,mRNA_file))
    for seq in SeqIO.parse(mRNA_file+'/'+i+'.fa','fasta'):
        if str(seq.name) == tblastn_item[i][1]:
            #print(tblastn_item[i])
            se = seq.seq
            # 获得orf
            #print(tblastn_item[i])
            s = int(tblastn_item[i][8])
            e = int(tblastn_item[i][9])
            if s>e:

                d = '-'
                se = se.reverse_complement()
            else:
                d = '+'
            #print(se)
            

            if d =='-':
                f =  (len(se) - e) %3 + 1
            else:
                f = s % 3 + 2
            
            se = se[f:]
            se = se[:len(se)-len(se)%3]
            #print(len(se)%3)

            frame = find_orfs_with_trans(se, 11, min_size)
            if len(frame) > 0:
                Mlocation = frame[0][0] + frame[0][2].find('M')*3
                mRNA = se[Mlocation:frame[0][1]]
                pep = mRNA.translate()

                outfile_mRNA.write('>'+i+'_mRNA\n'+str(mRNA)+'\n')
                
                outfile_pep.write('>'+i+'_pep\n'+str(pep)+'\n')
                print('spe %s has spe, lenth = %s'%(i,len(pep)) )
            else:
                print('%s 没找到正常的序列' % i)
            # start = frame[0][0]
            # end = frame[0][1]
            # mRNA = se[start:end]
            # pep = se.translate()
            # print(mRNA,pep)
outfile_pep.close()
outfile_mRNA.close()







    